import multiprocessing 
import time
import datetime
import numpy as np 

SHAPE = (2, 3)
global_arr_shared_1 = None
global_arr_shared_2 = None

def init_pool(arr_shared_1, arr_shared_2):
    global global_arr_shared_1
    global_arr_shared_1 = arr_shared_1
    global global_arr_shared_2
    global_arr_shared_2 = arr_shared_2

def worker(i):
    if i < 3:
        arr = np.frombuffer(global_arr_shared_1, np.double).reshape(SHAPE)
    if i > 2:
        arr = np.frombuffer(global_arr_shared_2, np.double).reshape(SHAPE)
    time.sleep(1)  # some other operations
    return np.sum(arr * i)


if __name__ == '__main__':
    #arr_1 = np.random.randn(*SHAPE)
    arr_1 = np.zeros(SHAPE, dtype = np.float32)
    arr_shared_1 = multiprocessing.RawArray('d', arr_1.ravel())
    #arr_2 = np.random.randn(*SHAPE)
    arr_2 = np.ones(SHAPE, dtype = np.float32)
    arr_shared_2 = multiprocessing.RawArray('d', arr_2.ravel())
    
    start_t = datetime.datetime.now()
    with multiprocessing.Pool(processes=4, initializer=init_pool, initargs=(arr_shared_1,arr_shared_2,)) as pool:  # initargs传入tuple
        result_1 = pool.apply_async(worker, args=(1,))
        result_2 = pool.apply_async(worker, args=(2,))
        result_3 = pool.apply_async(worker, args=(3,))
        result_4 = pool.apply_async(worker, args=(4,))
        print([result_1.get(), result_2.get(), result_3.get(), result_4.get()])
    end_t = datetime.datetime.now()
    print('Whole parallel time: %lf' % (end_t - start_t).total_seconds())
    
    start_t = datetime.datetime.now()
    global_arr_shared_1 = arr_shared_1
    global_arr_shared_2 = arr_shared_2
    for i in range(1,5):
        result = worker(i)
        print(result)
    end_t = datetime.datetime.now()
    print('Whole serial time: %lf' % (end_t - start_t).total_seconds())
